#pragma once

#include <string>
#include <map>
#include <set>
#include <fstream>
#include <sstream>
#include <vector>
#include <memory>

class Cd {
    class BaseAST {
    public:
        bool left_mode = false;
        virtual ~BaseAST() {}
        virtual void codegen() = 0;
    };

    class WhileAST : public BaseAST {
        std::shared_ptr<BaseAST> condition;
        std::shared_ptr<BaseAST> body;

    public:
        WhileAST(std::shared_ptr<BaseAST> condition, std::shared_ptr<BaseAST> body)
            : condition(condition), body(body) {}
        void codegen() override;
    };

    class ForAST : public BaseAST {
        std::string var_name;
        std::shared_ptr<BaseAST> start_expr;
        std::shared_ptr<BaseAST> end_expr;
        std::shared_ptr<BaseAST> step_expr;
    
        std::shared_ptr<BaseAST> body;

    public:
        ForAST(std::string var_name, std::shared_ptr<BaseAST> start_expr, std::shared_ptr<BaseAST> end_expr, std::shared_ptr<BaseAST> step_expr, std::shared_ptr<BaseAST> body)
            : var_name(var_name), start_expr(start_expr), end_expr(end_expr), step_expr(step_expr), body(body) {}
        void codegen() override;
    };

    class BreakAST : public BaseAST {
    public:
        void codegen() override;
    };

    class ContinueAST : public BaseAST {
    public:
        void codegen() override;
    };

    class IfAST : public BaseAST {
        std::shared_ptr<BaseAST> condition;
        std::shared_ptr<BaseAST> if_body;
        std::shared_ptr<BaseAST> else_body;

    public:
        IfAST(std::shared_ptr<BaseAST> condition, std::shared_ptr<BaseAST> if_body, std::shared_ptr<BaseAST> else_body)
            : condition(condition), if_body(if_body), else_body(else_body) {}
        void codegen() override;
    };

    class FnAST : public BaseAST {
        std::string fn_name;
        std::vector<std::string> arg_names;
        std::shared_ptr<BaseAST> body;

    public:
        FnAST(std::string fn_name, std::vector<std::string> arg_names, std::shared_ptr<BaseAST> body)
            : fn_name(fn_name), arg_names(arg_names), body(body) {}
        void codegen() override;
    };

    class RetAST : public BaseAST {
        std::shared_ptr<BaseAST> returnValue;

    public:
        RetAST(std::shared_ptr<BaseAST> returnValue)
            : returnValue(returnValue) {}
        void codegen() override;
    };

    class AsmAST : public BaseAST {
        std::string body;

    public:
        AsmAST(std::string body)
            : body(body) {}
        void codegen() override;
    };


    class StringAST : public BaseAST {
        std::string string;

    public:
        StringAST(std::string string)
            : string(string) {}
        void codegen() override;
    };

    class NumberAST : public BaseAST {
        double number;

    public:
        NumberAST(double number)
            : number(number) {}
        void codegen() override;
    };

    class VarAST : public BaseAST {
        std::string var_name;


    public:
        VarAST(std::string var_name)
            : var_name(var_name) {}
        void codegen() override;
    };

    class TableAST : public BaseAST {
        std::map<std::string, std::shared_ptr<BaseAST>> table;

    public:
        TableAST(std::map<std::string, std::shared_ptr<BaseAST>> table)
            : table(table) {}
        void codegen() override;
    };

    class CalcAST : public BaseAST {
        std::string op;
        bool is_union;
        std::shared_ptr<BaseAST> num1;
        std::shared_ptr<BaseAST> num2;

    public:
        CalcAST(std::string op, std::shared_ptr<BaseAST> num1, std::shared_ptr<BaseAST> num2, bool is_union = false)
            : op(op), num1(num1), num2(num2), is_union(is_union) {}
        void codegen() override;
    };

    class CallAST : public BaseAST {
        std::string fn_name;
        std::vector<std::shared_ptr<BaseAST>> args;

    public:
        CallAST(std::string fn_name, std::vector<std::shared_ptr<BaseAST>> args)
            : fn_name(fn_name), args(args) {}
        void codegen() override;
    };

    class AssignAST : public BaseAST {
        std::shared_ptr<BaseAST> left;
        std::shared_ptr<BaseAST> value;

    public:
        AssignAST(std::shared_ptr<BaseAST> left, std::shared_ptr<BaseAST> value)
            : left(left), value(value) {}
        void codegen() override;
    };

    class GetAST : public BaseAST {
    public:
        std::shared_ptr<BaseAST> table;
        std::shared_ptr<BaseAST> key;

        GetAST(std::shared_ptr<BaseAST> table, std::shared_ptr<BaseAST> key)
            : table(table), key(key) {}
        void codegen() override;
    };

    class BlockAST : public BaseAST {
        std::vector<std::shared_ptr<BaseAST>> exprs;

    public:
        BlockAST(std::vector<std::shared_ptr<BaseAST>> exprs)
            : exprs(exprs) {}
        void codegen() override;
    };

    enum TOKEN {
        TOKEN_WHILE,
        TOKEN_FOR,
        TOKEN_BREAK,
        TOKEN_CONTINUE,

        TOKEN_IF,
        TOKEN_ELSE,

        TOKEN_FN,
        TOKEN_RET,

        TOKEN_ASM,

        TOKEN_IDENT,
        TOKEN_NUM,
        TOKEN_STRING,
        TOKEN_SYMBOL,

        TOKEN_EOF,
    };

    std::map<std::string, int> op_table;
    std::set<std::string> symbols;

    static std::map<std::string, std::vector<std::string>> fns;

    std::string src;
    int index;

    int token = -1;
    std::string token_val = "";
    static std::vector<std::pair<std::string, std::string>> loop_tables;   // 记录循环的开始和结束标签名用于break和continue的汇编生成

    static std::set<std::string> buildin_fns;

    void next_token();
    void eat_token(std::string eat);

    std::shared_ptr<BaseAST> parse_expr();
    std::shared_ptr<BaseAST> parse_value();

public:
    static std::ofstream output_file;
    std::shared_ptr<BlockAST> parser();
    void init(std::string input_filename, std::string output_filename);
};